import matplotlib.pyplot as plt
import os

# 数据：对抗训练前的攻击成功率
before_training_success_rates = {
    'FGSM': 82.64,
    'PGD': 84.16,
    'CW': 97.61,
    'DeepFool': 96.29,
    'BIM': 5.69
}

# 数据：对抗训练后的攻击成功率
after_training_success_rates = {
    'FGSM': 40.25,
    'PGD': 43.79,
    'CW': 86.97,
    'DeepFool': 78.93,
    'BIM': 7.51
}

# 可视化函数
def plot_attack_success_rates(before, after, save_path):
    attack_names = list(before.keys())
    before_rates = [before[name] for name in attack_names]
    after_rates = [after[name] for name in attack_names]

    x = range(len(attack_names))  # 横坐标位置

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(x, before_rates, width=0.4, label='Before Training', align='center')
    ax.bar([i + 0.4 for i in x], after_rates, width=0.4, label='After PGD Training', align='center')

    ax.set_xlabel('Attack Methods')
    ax.set_ylabel('Attack Success Rate (%)')
    ax.set_title('Attack Success Rate')
    ax.set_xticks([i + 0.2 for i in x])
    ax.set_xticklabels(attack_names)
    ax.legend()

    plt.tight_layout()
    
    # 保存图片
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    plt.savefig(save_path)
    print(f"图片已保存到: {save_path}")

# 调用可视化函数，并保存图像
plot_attack_success_rates(before_training_success_rates, after_training_success_rates, save_path="./output/attack_success_rates.png")
